#include "string.h"
#include "yolox_post_process.h"

#include <sys/time.h>
#include <algorithm>
#include <cmath>
#include <future>
#include <iomanip>
#include <iostream>
#include <string>
#include <utility>
#include <vector>
#include "typeConv.hpp"
#include "common.h"
#ifdef LYNXI_PLUGIN
#include "lyn_plugin_dev.h"
#endif

using namespace COMMON;

struct Yolo8Config {
  std::vector<int> strides;
  std::vector<std::vector<std::pair<double, double>>> anchors_table;
  int class_num;
  std::vector<std::string> class_names;
};

enum class ModelType { Normal, X6 };

#define CLASS_NUM 2  // 80 //标准 yolov5 的分类为80，后处理在APU内部对应分类为2

static const Yolo8Config default_yolo8_config = {
    {8, 16, 32},
    {{{10, 13}, {16, 30}, {33, 23}},
     {{30, 61}, {62, 45}, {59, 119}},
     {{116, 90}, {156, 198}, {373, 326}}},
    CLASS_NUM,
    {"person",        "bicycle",      "car",
     "motorcycle",    "airplane",     "bus",
     "train",         "truck",        "boat",
     "traffic light", "fire hydrant", "stop sign",
     "parking meter", "bench",        "bird",
     "cat",           "dog",          "horse",
     "sheep",         "cow",          "elephant",
     "bear",          "zebra",        "giraffe",
     "backpack",      "umbrella",     "handbag",
     "tie",           "suitcase",     "frisbee",
     "skis",          "snowboard",    "sports ball",
     "kite",          "baseball bat", "baseball glove",
     "skateboard",    "surfboard",    "tennis racket",
     "bottle",        "wine glass",   "cup",
     "fork",          "knife",        "spoon",
     "bowl",          "banana",       "apple",
     "sandwich",      "orange",       "broccoli",
     "carrot",        "hot dog",      "pizza",
     "donut",         "cake",         "chair",
     "couch",         "potted plant", "bed",
     "dining table",  "toilet",       "tv",
     "laptop",        "mouse",        "remote",
     "keyboard",      "cell phone",   "microwave",
     "oven",          "toaster",      "sink",
     "refrigerator",  "book",         "clock",
     "vase",          "scissors",     "teddy bear",
     "hair drier",    "toothbrush"}};
template <class ForwardIterator>
inline size_t argmin(ForwardIterator first, ForwardIterator last) {
  return std::distance(first, std::min_element(first, last));
}

template <class ForwardIterator>
inline size_t argmax(ForwardIterator first, ForwardIterator last) {
  return std::distance(first, std::max_element(first, last));
}

typedef struct Bbox {
  float xmin;
  float ymin;
  float xmax;
  float ymax;

  Bbox() {}

  Bbox(float xmin, float ymin, float xmax, float ymax)
      : xmin(xmin), ymin(ymin), xmax(xmax), ymax(ymax) {}

  ~Bbox() {}
} Bbox;

typedef struct Detection {
  int id;
  float score;
  Bbox bbox;
  const char* class_name;
  Detection() {}

  Detection(int id, float score, Bbox bbox)
      : id(id), score(score), bbox(bbox) {}

  Detection(int id, float score, Bbox bbox, const char* class_name)
      : id(id), score(score), bbox(bbox), class_name(class_name) {}

  friend bool operator>(const Detection& lhs, const Detection& rhs) {
    return (lhs.score > rhs.score);
  }

  ~Detection() {}
} Detection;

void yolo8Nms(std::vector<Detection>& input, float iou_threshold, int top_k,
              std::vector<Detection>& result, bool suppress) {
  std::stable_sort(input.begin(), input.end(), std::greater<Detection>());

  std::vector<bool> skip(input.size(), false);

  std::vector<float> areas;
  areas.reserve(input.size());
  for (size_t i = 0; i < input.size(); i++) {
    float width = input[i].bbox.xmax - input[i].bbox.xmin;
    float height = input[i].bbox.ymax - input[i].bbox.ymin;
    areas.push_back(width * height);
  }

  int count = 0;
  for (size_t i = 0; /*count < top_k && */ i < skip.size(); i++) {
    if (skip[i]) {
      continue;
    }
    skip[i] = true;
    ++count;

    for (size_t j = i + 1; j < skip.size(); ++j) {
      if (skip[j]) {
        continue;
      }
      if (suppress == false) {
        if (input[i].id != input[j].id) {
          continue;
        }
      }

      float xx1 = std::max(input[i].bbox.xmin, input[j].bbox.xmin);
      float yy1 = std::max(input[i].bbox.ymin, input[j].bbox.ymin);
      float xx2 = std::min(input[i].bbox.xmax, input[j].bbox.xmax);
      float yy2 = std::min(input[i].bbox.ymax, input[j].bbox.ymax);

      if (xx2 > xx1 && yy2 > yy1) {
        float area_intersection = (xx2 - xx1) * (yy2 - yy1);
        float iou_ratio =
            area_intersection / (areas[j] + areas[i] - area_intersection);
        if (iou_ratio > iou_threshold) {
          skip[j] = true;
        }
      }
    }
    result.push_back(input[i]);
  }
}

const int sig[2] = {1, -1};
const float result = 5.96046e-08;

void tensorpostProcess(const Yolo8Config &yolo8_config, void* tensor, YoloxPostProcessInfo_t* post_info,
                       int layer, std::vector<Detection>& dets,
                       ModelType modelType) {
  // auto *data = reinterpret_cast<uint16_t *>(tensor);
  void* data = tensor;
  int stride = yolo8_config.strides[layer];
  int num_pred = post_info->class_num + 4;

  std::vector<int16_t> class_pred(post_info->class_num, 0);

  // std::vector<std::pair<double, double>>& anchors =
  //     yolo8_config.anchors_table[layer];
  double h_ratio = post_info->height * 1.0 / post_info->ori_height;
  double w_ratio = post_info->width * 1.0 / post_info->ori_width;
  double resize_ratio = std::min(w_ratio, h_ratio);
  if (post_info->is_pad_resize) {
    w_ratio = resize_ratio;
    h_ratio = resize_ratio;
  }

  int grid_height, grid_width;
  grid_height = post_info->height / stride;
  grid_width = post_info->width / stride;

  int16_t box_score_threshold = float2half(post_info->score_threshold);
  for (int h = 0; h < grid_height; h++) {
    for (int w = 0; w < grid_width; w++) {
      for (size_t k = 0; k < 1; k++) {
        int16_t* cur_data = (int16_t*)data + k * num_pred;
#if 0
                int16_t objness = cur_data[4];
                if (objness < box_score_threshold /*post_info->score_threshold*/) {
                    continue;
                }
                int32_t id = -1;
                double confidence = 0.0;
                if (modelType == ModelType::X7) {
                    id = cur_data[5];
                    confidence = half2float(objness) * half2float(cur_data[6]);
                } else if (modelType == ModelType::Normal) {
                    for (uint32_t index = 0; index < post_info->class_num; ++index) {
                        class_pred[index] = (cur_data[5 + index]);
                    }
                    id = argmax(class_pred.begin(), class_pred.end());
                    confidence = half2float(objness) * half2float(class_pred[id]);
                }
#else
        int32_t id = -1;
        double confidence = 0.0;
        int16_t objness = 0;
        if (modelType == ModelType::X6) {
          id = cur_data[4];
          objness = cur_data[5];
        } else if (modelType == ModelType::Normal) {
          for (int32_t index = 0; index < post_info->class_num; ++index) {
            class_pred[index] = (cur_data[4 + index]);
          }
          id = argmax(class_pred.begin(), class_pred.end());
          objness = class_pred[id];
        }
        if (objness < box_score_threshold) {
          continue;
        }
        confidence = half2float(objness);
#endif

        float center_x = half2float(cur_data[0]);
        float center_y = half2float(cur_data[1]);
        float scale_x = half2float(cur_data[2]);
        float scale_y = half2float(cur_data[3]);

        double xmin = (center_x - scale_x / 2.0);
        double ymin = (center_y - scale_y / 2.0);
        double xmax = (center_x + scale_x / 2.0);
        double ymax = (center_y + scale_y / 2.0);
        double w_padding =
            (post_info->width - w_ratio * post_info->ori_width) / 2.0;
        double h_padding =
            (post_info->height - h_ratio * post_info->ori_height) / 2.0;

        double xmin_org = (xmin - w_padding) / w_ratio;
        double xmax_org = (xmax - w_padding) / w_ratio;
        double ymin_org = (ymin - h_padding) / h_ratio;
        double ymax_org = (ymax - h_padding) / h_ratio;

        if (xmax_org <= 0 || ymax_org <= 0) {
          continue;
        }

        if (xmin_org > xmax_org || ymin_org > ymax_org) {
          continue;
        }

        xmin_org = std::max(xmin_org, 0.0);
        xmax_org = std::min(xmax_org, post_info->ori_width - 1.0);
        ymin_org = std::max(ymin_org, 0.0);
        ymax_org = std::min(ymax_org, post_info->ori_height - 1.0);

        Bbox bbox(xmin_org, ymin_org, xmax_org, ymax_org);
        int classNameSize = yolo8_config.class_names.size();
        if (id > classNameSize || id < 0) {
          // LOG_PLUGIN_E("ERROR!!! id: %d\n", id);
          continue;
        }
        dets.push_back(
            Detection((int)id, confidence, bbox,
                      yolo8_config.class_names[(int)id].c_str()));
      }
      data = (int16_t*)data + num_pred * 1;
    }
  }
}

int lynYoloxPostProcess(YoloxPostProcessInfo_t* post_info) {
  int16_t* output_tensor = nullptr;
  lynBoxesInfo* boxesInfo = nullptr;
#ifdef LYNXI_PLUGIN
  output_tensor = (int16_t*)lynPluginGetVirtAddr(post_info->output_tensor);
  if (output_tensor == nullptr) {
    LOG_PLUGIN_E("get output tensor addr error\n");
    return -1;
  }

  boxesInfo = (lynBoxesInfo*)lynPluginGetVirtAddr(post_info->boxesInfo);
  if (boxesInfo == nullptr) {
    LOG_PLUGIN_E("get boxesInfo addr error\n");
    return -1;
  }
#else
  output_tensor = (int16_t*)post_info->output_tensor;
  boxesInfo = post_info->boxesInfo;
#endif

  Yolo8Config yolo8_config = default_yolo8_config;
  generateClassNameAndNum(post_info->labelList,
                          yolo8_config.class_names,
                          yolo8_config.class_num);

  ModelType modelType = ModelType::X6;
  void* tensor0 = NULL;
  tensor0 = output_tensor;
  void* tensor1 = NULL;
  int grid_width1 = post_info->width / yolo8_config.strides[0];
  int grid_height1 = post_info->height / yolo8_config.strides[0];
  tensor1 = (int16_t*)tensor0 +
            1 * 1 * grid_width1 * grid_height1 * (post_info->class_num + 4);
  void* tensor2 = NULL;
  int grid_width2 = post_info->width / yolo8_config.strides[1];
  int grid_height2 = post_info->height / yolo8_config.strides[1];
  tensor2 = (int16_t*)tensor1 +
            1 * 1 * grid_width2 * grid_height2 * (post_info->class_num + 4);
  std::vector<Detection> dets;
  std::vector<Detection> det_restuls;
  // size_t i = 0;
  tensorpostProcess(yolo8_config, tensor0, post_info, 0, dets, modelType);
  tensorpostProcess(yolo8_config, tensor1, post_info, 1, dets, modelType);
  tensorpostProcess(yolo8_config, tensor2, post_info, 2, dets, modelType);
  yolo8Nms(dets, post_info->nms_threshold, post_info->nms_top_k, det_restuls,
           false);

  boxesInfo->boxesNum = std::min((int)det_restuls.size(), BOX_MAX_NUM);

  for (int i = 0; i < boxesInfo->boxesNum; i++) {
    boxesInfo->boxes[i].xmax = det_restuls[i].bbox.xmax;
    boxesInfo->boxes[i].xmin = det_restuls[i].bbox.xmin;
    boxesInfo->boxes[i].ymax = det_restuls[i].bbox.ymax;
    boxesInfo->boxes[i].ymin = det_restuls[i].bbox.ymin;
    boxesInfo->boxes[i].score = det_restuls[i].score;
    boxesInfo->boxes[i].id = det_restuls[i].id;
    strncpy(boxesInfo->boxes[i].label, det_restuls[i].class_name,
            sizeof(boxesInfo->boxes[i].label) - 1);
    if (static_cast<int>(boxesInfo->boxes[i].xmax) > post_info->ori_width - 1)
      boxesInfo->boxes[i].xmax = post_info->ori_width - 1;
    if (static_cast<int>(boxesInfo->boxes[i].xmin) < 0)
      boxesInfo->boxes[i].xmin = 0;
    if (static_cast<int>(boxesInfo->boxes[i].ymax) > post_info->ori_height - 1)
      boxesInfo->boxes[i].ymax = post_info->ori_height - 1;
    if (static_cast<int>(boxesInfo->boxes[i].ymin) < 0)
      boxesInfo->boxes[i].ymin = 0;
  }

  return 0;
}